package dogapm

import (
	"context"
	"dogapm/internal"
	"fmt"
	"github.com/redis/go-redis/v9"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.24.0"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"gorm.io/driver/mysql"
	_ "gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"time"
)

const (
	ctxKeyBeginTime = "begin_time"
	mysqlTracerName = "dogapm/mysql"
)

type infra struct {
	RDB     *redis.Client
	MySqlDB *gorm.DB
}

var Infra = &infra{}

type InfraOption func(*infra)

func WithRDB(addr string) InfraOption {
	return func(i *infra) {
		client := redis.NewClient(&redis.Options{
			Addr: addr,
			DB:   0,
		})
		client.AddHook(&redisHook{})
		ping, err := client.Ping(context.Background()).Result()
		if err != nil {
			panic(err)
		}
		if ping != "PONG" {
			panic("redis connect failed")
		}
		i.RDB = client
	}
}

func WithMySql(addr string) InfraOption {
	return func(i *infra) {
		g, _ := gorm.Open(mysql.New(mysql.Config{
			DSN:                       fmt.Sprintf("%s?charset=utf8mb4&parseTime=True&loc=Local", addr), // data source name
			DefaultStringSize:         256,                                                              // default size for string fields
			DisableDatetimePrecision:  true,                                                             // disable datetime precision, which not supported before MySQL 5.6
			DontSupportRenameIndex:    true,                                                             // drop & create when rename index, rename index not supported before MySQL 5.7, MariaDB
			DontSupportRenameColumn:   true,                                                             // `change` when rename column, rename column not supported before MySQL 8, MariaDB
			SkipInitializeWithVersion: false,
		}), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Info),
		})

		var gCtx context.Context
		_ = g.Callback().Query().Before("gorm:query").Register("gorm:before_query_for_benchmark", func(db *gorm.DB) {
			gCtx = db.Statement.Context
			db.InstanceSet("start_time", time.Now()) // 设置开始时间
			tracer := otel.Tracer(mysqlTracerName)
			if ctx, span := tracer.Start(gCtx, "sqltrace"); span != nil {
				span.SetAttributes(
					attribute.String("sql", db.Statement.SQL.String()),
				)
				gCtx = ctx
			}
		})

		_ = g.Callback().Query().After("gorm:query").Register("gorm:after_query_for_benchmark", func(db *gorm.DB) {
			startTime, ok := db.InstanceGet("start_time")
			if !ok {
				return
			}

			now := time.Now()
			span := trace.SpanFromContext(gCtx)
			span.SetAttributes(
				attribute.Bool("slowsql", true),
			)
			if now.Sub(startTime.(time.Time)).Seconds() >= 1 {
				span.SetAttributes(
					attribute.Bool("slowsql", true),
				)
			}
			span.End()
		})
		i.MySqlDB = g
	}
}

func (i *infra) NewInfra(opts ...InfraOption) {
	for _, opt := range opts {
		opt(i)
	}
	Tracer = otel.Tracer("dogapm")
}

func InfraEnableAPM(
	otelEndpoint string) InfraOption {
	return func(inf *infra) {
		ctx := context.Background()
		res, err := resource.New(ctx,
			resource.WithAttributes(
				semconv.ServiceName(internal.BuildInfo.AppName()),
			),
		)
		if err != nil {
			panic(err)
		}

		ctx, cancel := context.WithTimeout(ctx, time.Second)
		defer cancel()
		conn, err := grpc.NewClient(otelEndpoint,
			grpc.WithTransportCredentials(insecure.NewCredentials()),
		)
		if err != nil {
			panic(err)
		}
		// Set up a trace exporter
		traceExporter, err := otlptracegrpc.New(ctx, otlptracegrpc.WithGRPCConn(conn))
		if err != nil {
			panic(err)
		}
		bsp := sdktrace.NewBatchSpanProcessor(traceExporter)
		tracerProvider := sdktrace.NewTracerProvider(
			sdktrace.WithSampler(sdktrace.AlwaysSample()),
			sdktrace.WithResource(res),
			sdktrace.WithSpanProcessor(bsp),
		)
		otel.SetTracerProvider(tracerProvider)
		otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))
		globalClosers = append(globalClosers, &traceProviderComponent{provider: tracerProvider})
	}
}

type traceProviderComponent struct {
	provider *sdktrace.TracerProvider
}

func (t traceProviderComponent) Close() {
	_ = t.provider.Shutdown(context.TODO())
}
